import numpy as np
import trimesh
from trimesh.voxel.creation import voxelize


def R_matrix(z_angle,y_angle,x_angle):
    # s angle: row; y angle: pitch; z angle: yaw
    # generate rotation matrix in SO3
    # RzRyRx = R, ZYX intrinsic rotation
    # also (r1,r2,r3) in R3*3 in {W} frame
    # used in obb.O
    # [[R p]
    # [0T 1]] gives transformation from body to world
    return np.array([[np.cos(z_angle), -np.sin(z_angle), 0.0], [np.sin(z_angle), np.cos(z_angle), 0.0], [0.0, 0.0, 1.0]])@ \
           np.array([[np.cos(y_angle), 0.0, np.sin(y_angle)], [0.0, 1.0, 0.0], [-np.sin(y_angle), 0.0, np.cos(y_angle)]])@ \
           np.array([[1.0, 0.0, 0.0], [0.0, np.cos(x_angle), -np.sin(x_angle)], [0.0, np.sin(x_angle), np.cos(x_angle)]])

def getblocks():
    # AABBs
    block = [[4.00e+00, 1.20e+01, 0.00e+00, 5.00e+00, 2.00e+01, 5.00e+00],
             [5.5e+00, 1.20e+01, 0.00e+00, 1.00e+01, 1.30e+01, 5.00e+00],
             [1.00e+01, 1.20e+01, 0.00e+00, 1.40e+01, 1.30e+01, 5.00e+00],
             [1.00e+01, 9.00e+00, 0.00e+00, 2.00e+01, 1.00e+01, 5.00e+00],
             [9.00e+00, 6.00e+00, 0.00e+00, 1.00e+01, 1.00e+01, 5.00e+00]]
    Obstacles = []
    for i in block:
        i = np.array(i)
        Obstacles.append([j for j in i])
    return np.array(Obstacles)

def getballs():
    spheres = [[2.0,6.0,2.5,1.0],[14.0,14.0,2.5,2]]
    Obstacles = []
    for i in spheres:
        Obstacles.append([j for j in i])
    return np.array(Obstacles)

def getAABB(blocks):
    # used for Pyrr package for detecting collision
    AABB = []
    for i in blocks:
        AABB.append(np.array([np.add(i[0:3], -0), np.add(i[3:6], 0)]))  # make AABBs alittle bit of larger
    return AABB

def getAABB2(blocks):
    # used in lineAABB
    AABB = []
    for i in blocks:
        AABB.append(aabb(i))
    return AABB

def add_block(block = [1.51e+01, 0.00e+00, 2.10e+00, 1.59e+01, 5.00e+00, 6.00e+00]):
    return block


class aabb(object):
    # make AABB out of blocks,
    # P: center point
    # E: extents
    # O: Rotation matrix in SO(3), in {w}
    def __init__(self,AABB):
        self.P = [(AABB[3] + AABB[0])/2, (AABB[4] + AABB[1])/2, (AABB[5] + AABB[2])/2]# center point
        self.E = [(AABB[3] - AABB[0])/2, (AABB[4] - AABB[1])/2, (AABB[5] - AABB[2])/2]# extents
        self.O = [[1,0,0],[0,1,0],[0,0,1]]


class obb(object):
    # P: center point
    # E: extents
    # O: Rotation matrix in SO(3), in {w}
    def __init__(self, P, E, O):
        self.P = P
        self.E = E
        self.O = O
        self.T = np.vstack([np.column_stack([self.O.T,-self.O.T@self.P]),[0,0,0,1]])


class Maze():
    def __init__(self, map, start=[0,0,0], target=[0,0,5], resolution=1):
        self.map = map
        self.start = start
        self.target = target
        self.resolution = resolution
        #self.boundary = np.array([xmin, ymin, zmin, xmax, ymax, zmax])
        self.boundary = np.hstack([self.map.bounds[0], self.map.bounds[1]])
        self.blocks = np.array([[0, 0, 0, 0, 0, 0]])
        self.AABB = getAABB2(self.blocks)
        self.AABB_pyrr = getAABB(self.blocks)
        self.balls = []
        self.OBB = np.array([])
        self.t = 0

    def add_balls(self, spheres):
        # spheres = [[2.0, 6.0, 2.5, 1.0], [14.0, 14.0, 2.5, 2]]
        Obstacles = []
        for i in spheres:
            Obstacles.append([j for j in i])
        self.balls = np.array(Obstacles)

    def add_OBB(self, obb):
        self.OBB=np.append(self.OBB, obb)

    def add_blocks(self, blocks):
        Obstacles = []
        for i in blocks:
            i = np.array(i)
            Obstacles.append([j for j in i])
        self.blocks = np.array(Obstacles)

    def set_start(self, start):
        self.start = start

    def set_target(self, target):
        self.target = target

    def set_boundaries(self, bounds):
        if self.map is not None:
            if bounds[0] is not None:
                self.map = self.map.slice_plane(bounds[0], [1, 0, 0])
                self.map = self.map.slice_plane(bounds[0], [0, 1, 0])
                self.map = self.map.slice_plane(bounds[0], [0, 0, 1])
            if bounds[1] is not None:
                self.map = self.map.slice_plane(bounds[1], [-1, 0, 0])
                self.map = self.map.slice_plane(bounds[1], [0, -1, 0])
                self.map = self.map.slice_plane(bounds[1], [0, 0, -1])

        return self.map

    def get_bounds(self):
        return self.map.bounds

#     def get_pointcloud(self, size=64):
# #        max_side = (self.map.bounds[1] - self.map.bounds[0]).max()
# #        vox = voxelize(self.map, max_side / size)
#         vox = voxelize(self.map, self.resolution)
#         pc = trimesh.PointCloud(vox.points)
#         return pc

    def get_pointcloud(self):
        vox = voxelize(self.map, self.resolution)
        pc = trimesh.PointCloud(vox.points)
        return pc

    def add_element(self):
        pass

    def voxelized(self):
        #max_side = (self.map.bounds[1] - self.map.bounds[0]).max()
        vox = voxelize(self.map, self.resolution)
        matrix = vox.matrix
        return matrix

    def solve(self):
        pass

